from py2neo import Graph, Node, Relationship, NodeMatcher

from settings import (
  NEO4J_URL,
  NEO4J_USERNAME,
  NEO4J_PASSWORD,
)

from utils import (
  log_message
)

from consts import (
  NODE_MOVIE
)

# 连接函数
def connect():
  try:
    graph = Graph(NEO4J_URL, auth=(NEO4J_USERNAME, NEO4J_PASSWORD), name='neo4j')
    log_message('数据库连接成功')
    return graph
  except:
    log_message('数据库连接失败')

graph = connect()
    
# 创建节点
def create_one_node(node_type, attrs_obj):
  node = Node(node_type, **attrs_obj)
  graph.create(node)

# 创建多个节点
def create_multiple_node(node_type, attrs_list):
  # 事务开始
  tx = graph.begin()
  total = len(attrs_list)
  count = 0
  for attrs in attrs_list:
    count += 1
    if type(attrs) == dict:
      node = Node(node_type, **attrs)
    if type(attrs) == str:
      node = Node(node_type, name=attrs)

    tx.create(node)
    if count % 800 == 0:
      log_message('存储%s节点数据进度：%.2f%%' % (node_type, count / total * 100))
      graph.commit(tx)
      tx = graph.begin()
  # 执行
  graph.commit(tx)
  log_message(f'存储{node_type}节点完成')

# 创建实体关系
def create_relationship(start_node, relationship_type , end_node, relationship_attrs=None):
  if relationship_attrs != None:
    relation = Relationship(start_node, relationship_type, end_node, **relationship_attrs)
  else:
    relation = Relationship(start_node, relationship_type, end_node)
  graph.create(relation)



# 创建多个相同类型的实体关系,头节点不变，尾节点多个
def create_multiple_relationship_one_to_more(relationship_type, node_list):
  '''
  node_list数据结构
  [
    {
      start_node:...,
      end_node_list:[],
      relationship_attrs:{
        ...
      },
    },
    ...
  ]
  '''
  tx = graph.begin()
  count = 0
  for item in node_list:
    count += 1
    start_node = item['start_node']
    end_node_list = item['end_node_list']
    relationship_attrs = item['relationship_attrs']
    for end_node in end_node_list:
      relation = Relationship(start_node, relationship_type, end_node, **relationship_attrs)
      tx.create(relation)
    if count % 100 == 0:
      graph.commit(tx)
      tx = graph.begin()

  graph.commit(tx)


def create_multiple_relationship(relationship_type, node_list):
  '''
  node_list数据结构
  [
    {
      start_node:...,
      end_node:...,
      relationship_attrs:{
        ...
      },
    },
    ...
  ]
  '''
  tx = graph.begin()
  count = 0
  for item in node_list:
    count += 1
    start_node = item['start_node']
    end_node = item['end_node']
    relationship_attrs = item['relationship_attrs']
    relation = Relationship(start_node, relationship_type, end_node, **relationship_attrs)
    tx.create(relation)
    if count % 100 == 0:
      graph.commit(tx)
      tx = graph.begin()
  graph.commit(tx)


# 查找某个节点
def find_node_by_attr(node_type, attr_key, attr_value):
  matcher = NodeMatcher(graph)
  if attr_key == 'id':
    node = matcher.match(node_type).where(f'_.{attr_key} = {attr_value}').first()
  else:
    node = matcher.match(node_type).where(f'_.{attr_key} = "{attr_value}"').first()
  # 这个id是内部Neo4j分配的节点ID
  # node = matcher.match(node_type, id=node_id).first()
  if not node:
    log_message(f'未找到{node_type}节点')
  return node

# 清除图数据库中所有的节点
def clear_graph():
  while True:
    # 获取剩余节点总数
    count_query = 'MATCH (n) RETURN count(n) AS total'
    total_nodes = graph.evaluate(count_query)
    # 检查是否还有节点需要删除
    if total_nodes == 0:
      break
    batch_size = 1000
    # 随机删除一个批次的节点
    query = f'MATCH (n) WITH n ORDER BY RAND() LIMIT {batch_size} DETACH DELETE n'
    graph.run(query)
  log_message('所有节点删除成功')

# 清除数据库中某种类型的节点
def clear_node_by_type(node_type, batch_size=1000):
  while True:
    # Cypher 查询：分批匹配和删除节点
    query = f"""
    MATCH (n:{node_type})
    WITH n LIMIT {batch_size}
    DETACH DELETE n
    RETURN count(n) as deleted_count
    """
    result = graph.run(query).data()

    # 如果删除的节点数量为 0，停止循环
    if result[0]['deleted_count'] == 0:
      break
  log_message(f"{node_type}实体关系删除完毕")

# 清楚图数据库中所有的实体关系
def clear_relationship():
  query = f'MATCH ()-[r]->() DELETE r'
  result = graph.run(query)
  log_message('所有关系删除完毕')

def delete_relationship_by_type(relationship_type):
  query = f'MATCH ()-[r:{relationship_type}]->() DELETE r'
  result = graph.run(query)
  log_message(f'{relationship_type}实体关系删除完毕')


# 测试
# relation = Relationship(start_node, relationship_type, end_node, **relationship_attrs)